package cn.edu.buaa.rwntm;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.IOException;
import java.util.*;

/**
 * Created by dell on 2016/4/20.
 */
public class RWNTM implements Runnable {
    private final String data_path;
    private final double alpha;
    private final double beta0;
    private final double beta1;
    private final double gama;
    private final int iter_num;
    private final int topic_num;
    private final String outputPath;

    public HashMap<String, Integer> wordMap = new HashMap<>();
    public int wordMapSize = 0;
    public HashMap<Integer, String> id2Word = new HashMap<>();

    public int[][] wordId_of_corpus = null;
    public int[][] topicId_of_corpus = null;

    public int[][] word_num_in_topic_word = null;
    public int[] word_num_in_topic = null;
    public int[][] word_num_in_doc_topic = null;
    public int[] word_num_in_doc_of_topic = null;

//    private int[] background_word = null;
//    private int background_word_sum = 0;

    private SparseIntMatrix doc_specific_word = null;
    private int[] doc_specific_word_sum = null;

    private int[][] namota = null;
    private int[][] word_switch = null;
    private List<List<String>> doc;


    public RWNTM(String data_path, int topic_num, int iter_num, double alpha, double beta, double gamma, int threadId) {

        this.alpha = alpha;
        this.beta0 = beta;
        this.beta1 = 0.0001;
//        this.beta2 = beta;
        this.gama = gamma;
        this.iter_num = iter_num;
        this.topic_num = topic_num;
        this.data_path = data_path;
        this.outputPath = (new File(data_path)).getParentFile().getAbsolutePath() + "/rwntm_" + alpha + "_" + beta + "_" + gamma + "_" + topic_num + "/" + threadId + "/";
        (new File(this.outputPath)).mkdirs();
        print_parameter();
    }

    public RWNTM(List<List<String>> doc, int topic_num) {

        this.alpha = 0.1;
        this.beta0 = 0.01;
        this.beta1 = 0.0001;
//        this.beta2 = beta;
        this.gama = 0.3;
        this.iter_num = 500;
        this.topic_num = topic_num;
        this.doc = doc;
        data_path = "";
        outputPath = "";
        print_parameter();


    }

    public static String build_RWNTM(String data_path, int topic_num, int iter_num, double alpha, double beta, double gamma, int threadNum) {

        RWNTM model = new RWNTM(data_path, topic_num, iter_num, alpha, beta, gamma, 0);
        model.run();
        return (new File(data_path)).getParentFile().getAbsolutePath() + "/rwntm_" + alpha + "_" + beta + "_" + gamma + "_" + topic_num + "/" + 0 + "/model-final.theta";

    }

    private void print_parameter() {
        System.out.println("method:GSLDA-without-background"
                + "	path:" + this.data_path
                + "\talpha:" + this.alpha
                + "\tbeta0:" + this.beta0
                + "\tbeta1:" + this.beta1
                + "\tgama:" + this.gama
                + "\titer_num:" + this.iter_num
                + "\ttopic_num:" + this.topic_num
        );
    }

    private void load_data() {
        try {
            BufferedReader reader = IOUtil.getReader(this.data_path);
            String line;
            ArrayList<int[]> tmpCorpus = new ArrayList<>();
            while ((line = reader.readLine()) != null) {
                String[] words = line.split("\\s+");
                int[] oneDoc = new int[words.length];
                for (int wordIndex = 0; wordIndex < words.length; wordIndex++) {
                    if (!this.wordMap.containsKey(words[wordIndex])) {
                        this.wordMap.put(words[wordIndex], this.wordMap.size());
                        this.id2Word.put(this.wordMap.get(words[wordIndex]), words[wordIndex]);
                    }
                    oneDoc[wordIndex] = this.wordMap.get(words[wordIndex]);
                }
                tmpCorpus.add(oneDoc);
            }
            reader.close();
            this.wordId_of_corpus = new int[tmpCorpus.size()][];
            for (int docIndex = 0; docIndex < this.wordId_of_corpus.length; docIndex++) {
                this.wordId_of_corpus[docIndex] = tmpCorpus.get(docIndex);
            }

        } catch (Exception e) {
            e.printStackTrace();
        }

    }

    private void load_data2() {

        ArrayList<int[]> tmpCorpus = new ArrayList<>();
        for (List<String> words : this.doc) {
            int[] oneDoc = new int[words.size()];
            for (int wordIndex = 0; wordIndex < words.size(); wordIndex++) {
                if (!this.wordMap.containsKey(words.get(wordIndex))) {
                    this.wordMap.put(words.get(wordIndex), this.wordMap.size());
                    this.id2Word.put(this.wordMap.get(words.get(wordIndex)), words.get(wordIndex));
                }
                oneDoc[wordIndex] = this.wordMap.get(words.get(wordIndex));
            }
            tmpCorpus.add(oneDoc);
        }
        this.wordId_of_corpus = new int[tmpCorpus.size()][];
        for (int docIndex = 0; docIndex < this.wordId_of_corpus.length; docIndex++) {
            this.wordId_of_corpus[docIndex] = tmpCorpus.get(docIndex);
        }

    }

    private void init_model() {

        this.wordMapSize = this.wordMap.size();

        this.topicId_of_corpus = new int[this.wordId_of_corpus.length][];

        this.word_num_in_topic_word = new int[this.topic_num][this.wordMap.size()];
        this.word_num_in_topic = new int[this.topic_num];
        this.word_num_in_doc_topic = new int[this.wordId_of_corpus.length][this.topic_num];
        this.word_num_in_doc_of_topic = new int[this.wordId_of_corpus.length];

//        this.background_word = new int[this.wordMap.size()];
//        this.background_word_sum = 0;
        this.doc_specific_word = new SparseIntMatrix(this.wordId_of_corpus.length, this.wordMap.size());
        this.doc_specific_word_sum = new int[this.wordId_of_corpus.length];

        this.namota = new int[this.wordId_of_corpus.length][2];
        this.word_switch = new int[this.wordId_of_corpus.length][];


        for (int docIndex = 0; docIndex < this.wordId_of_corpus.length; docIndex++) {
            this.topicId_of_corpus[docIndex] = new int[this.wordId_of_corpus[docIndex].length];
//            this.word_num_in_doc[docIndex] = this.wordId_of_corpus[docIndex].length;
            this.word_switch[docIndex] = new int[this.wordId_of_corpus[docIndex].length];


            for (int wordIndex = 0; wordIndex < this.wordId_of_corpus[docIndex].length; wordIndex++) {
                int wordId = this.wordId_of_corpus[docIndex][wordIndex];
                int switchId = (int) (Math.random() * 2);

                this.namota[docIndex][switchId] += 1;
                this.word_switch[docIndex][wordIndex] = switchId;

                switch (switchId) {
                    case 0:
                        int topicId = (int) (Math.random() * this.topic_num);
                        this.topicId_of_corpus[docIndex][wordIndex] = topicId;
                        this.word_num_in_topic_word[topicId][wordId] += 1;
                        this.word_num_in_topic[topicId] += 1;
                        this.word_num_in_doc_topic[docIndex][topicId] += 1;
                        this.word_num_in_doc_of_topic[docIndex] += 1;
                        break;
                    case 1:
                        this.doc_specific_word.add(docIndex, wordId, 1);
                        this.doc_specific_word_sum[docIndex] += 1;
                        break;
                    default:
//                        this.background_word[wordId] += 1;
//                        this.background_word_sum += 1;
                        assert switchId > -1;
                        assert switchId < 2;
                        break;
                }
            }
        }
    }

    private void save_twords(int topWordNum) throws IOException {
        BufferedWriter writer = IOUtil.getWriter(this.outputPath + "model-final.twords");

        for (int topic_index = 0; topic_index < this.word_num_in_topic_word.length; topic_index++) {
            HashMap<Integer, Double> oneLine = new HashMap<>();
            for (int word_id = 0; word_id < this.word_num_in_topic_word[topic_index].length; word_id++) {
                oneLine.put(word_id, ((double) this.word_num_in_topic_word[topic_index][word_id]) / this.word_num_in_topic[topic_index]);
            }
            List<Map.Entry<Integer, Double>> maplist = new ArrayList<>(oneLine.entrySet());

            Collections.sort(maplist, new Comparator<Map.Entry<Integer, Double>>() {
                public int compare(Map.Entry<Integer, Double> o1, Map.Entry<Integer, Double> o2) {
                    return (o2.getValue().compareTo(o1.getValue()));
                }
            });

            writer.write("Topic:" + topic_index + "\n");
            int count = 0;
            for (Map.Entry<Integer, Double> o1 : maplist) {
                writer.write("\t" + this.id2Word.get(o1.getKey()) + "\t:" + o1.getValue() + "\n");
                count++;
                if (count > topWordNum) {
                    break;
                }
            }
            writer.write("\n");

        }


        writer.close();
    }

    private void save_wordMap() throws IOException {
        BufferedWriter writer = IOUtil.getWriter(this.outputPath + "model-final.wordmap");
        for (String key : this.wordMap.keySet()) {
            writer.write(key + " " + this.wordMap.get(key));
            writer.write("\n");
        }
        writer.close();
    }

    private void save_theta() throws IOException {
        BufferedWriter writer = IOUtil.getWriter(this.outputPath + "model-final.theta");
        int doc_index = 0;
        for (int[] line : this.word_num_in_doc_topic) {
            for (int topic : line) {
                writer.write((this.alpha + topic) / (this.wordId_of_corpus[doc_index].length + this.alpha * this.topic_num) + " ");
//                writer.write(((double)topic) / (this.word_num_in_doc[doc_index]) + " ");
            }
            writer.write("\n");
            doc_index++;
        }
        writer.close();
    }

    private void save_tassign() throws IOException {
        BufferedWriter writer = IOUtil.getWriter(this.outputPath + "model-final.tassign");
        for (int doc_index = 0; doc_index < this.topicId_of_corpus.length; doc_index++) {
            for (int word_index = 0; word_index < this.topicId_of_corpus[doc_index].length; word_index++) {
                if (this.word_switch[doc_index][word_index] == 0) {
                    writer.write(this.wordId_of_corpus[doc_index][word_index] + ":" + this.topicId_of_corpus[doc_index][word_index] + " ");
                }
            }
            writer.write("\n");
        }
        writer.close();
    }

    private void save_phi() throws IOException {
        BufferedWriter writer = IOUtil.getWriter(this.outputPath + "model-final.phi");
        int topic_index = 0;
        for (int[] line : this.word_num_in_topic_word) {
            for (int words : line) {
                writer.write((this.beta0 + words) / (this.word_num_in_topic[topic_index] + this.beta0 * this.wordMap.size()) + " ");
            }
            writer.write("\n");
            topic_index++;
        }
        writer.close();

    }

//    private void save_background_words() throws IOException {
//        HashSet<Integer> unique = new HashSet<>();
//        BufferedWriter writer = IOUtil.getWriter(this.outputPath + "model-final.bgw");
//        HashMap<Integer, Double> oneLine = new HashMap<>();
//        for (int word_id = 0; word_id < this.background_word.length; word_id++) {
//            oneLine.put(word_id, ((double)this.background_word[word_id]) / this.background_word.length);
//        }
//        List<Map.Entry<Integer, Double>> maplist = new ArrayList<>(oneLine.entrySet());
//        Collections.sort(maplist, (o1, o2) -> o2.getValue().compareTo(o1.getValue()));
//        for(int i = 0; i< this.background_word.length;i++){
//            writer.write(id2Word.get(maplist.get(i).getKey())+":"+maplist.get(i).getValue()+"\n");
//        }
//        writer.close();
//    }

    private void save_specific_words() throws IOException {
        BufferedWriter writer = IOUtil.getWriter(this.outputPath + "model-final.spec");
        for (int doc_index = 0; doc_index < this.topicId_of_corpus.length; doc_index++) {
            for (int word_index = 0; word_index < this.topicId_of_corpus[doc_index].length; word_index++) {
                if (this.word_switch[doc_index][word_index] == 1) {
                    writer.write(this.wordId_of_corpus[doc_index][word_index] + ":" + this.id2Word.get(this.wordId_of_corpus[doc_index][word_index]) + " ");
                }
            }
            writer.write("\n");
        }
        writer.close();
    }

    private int get_pid(int docIndex, int wordId) {
        int pid = -1;
        int p_num = this.topic_num + 1;
        double[] p = new double[p_num];

        for (int k = 0; k < this.topic_num; k++) {
            p[k] = (this.namota[docIndex][0] + this.gama)
                    * (this.word_num_in_doc_topic[docIndex][k] + this.alpha)
                    * (this.word_num_in_topic_word[k][wordId] + this.beta0)
//                    / (this.wordId_of_corpus[docIndex].length - 1 + 2 * this.gama)
                    / (this.word_num_in_doc_of_topic[docIndex] + this.topic_num * this.alpha)
                    / (this.word_num_in_topic[k] + this.wordMapSize * this.beta0);
        }

        p[p_num - 1] = (this.namota[docIndex][1] + this.gama)
                * (this.doc_specific_word.get(docIndex, wordId) + this.beta1)
//                / (this.wordId_of_corpus[docIndex].length - 1 + 2 * this.gama)
                / (this.doc_specific_word_sum[docIndex] + this.wordMapSize * this.beta1);

//        p[p_num-1] = (this.namota[docIndex][2]                 + this.gama)
//                    *(this.background_word[wordId]             + this.beta2)
//                    /(this.wordId_of_corpus[docIndex].length-1 +                3 * this.gama)
//                    /(this.background_word_sum                 + this.wordMapSize * this.beta2);

        for (int k = 1; k < p_num; k++) {
            p[k] += p[k - 1];
        }

        double u = Math.random() * p[p_num - 1];
        for (int k = 0; k < p_num; k++) {
            if (u < p[k]) {
                pid = k;
                break;
            }

        }

        return pid;
    }

    private int sample(int docIndex, int wordIndex) {
        int wordId = this.wordId_of_corpus[docIndex][wordIndex];
        int oldSwitchId = this.word_switch[docIndex][wordIndex];

        int newTopicId = -1;
        int newSwitchId = -1;

        this.namota[docIndex][oldSwitchId] -= 1;

        switch (oldSwitchId) {
            case 0:
                int oldTopicId = this.topicId_of_corpus[docIndex][wordIndex];
                this.word_num_in_topic_word[oldTopicId][wordId] -= 1;
                this.word_num_in_topic[oldTopicId] -= 1;
                this.word_num_in_doc_topic[docIndex][oldTopicId] -= 1;
                this.word_num_in_doc_of_topic[docIndex] -= 1;
                break;
            case 1:
                this.doc_specific_word.add(docIndex, wordId, -1);
                this.doc_specific_word_sum[docIndex] -= 1;
                break;
            default:
                assert oldSwitchId > -1 && oldSwitchId < 2;
                break;
        }

        newTopicId = get_pid(docIndex, wordId);
        assert newTopicId > -1;
        if (newTopicId < this.topic_num) {
            newSwitchId = 0;
            this.word_num_in_topic_word[newTopicId][wordId] += 1;
            this.word_num_in_topic[newTopicId] += 1;
            this.word_num_in_doc_topic[docIndex][newTopicId] += 1;
            this.word_num_in_doc_of_topic[docIndex] += 1;

        } else if (newTopicId == this.topic_num) {
            newSwitchId = 1;
            this.doc_specific_word.add(docIndex, wordId, 1);
            this.doc_specific_word_sum[docIndex] += 1;

        } else {
            newSwitchId = -1;
            assert newTopicId < this.topic_num + 1;
        }

        this.namota[docIndex][newSwitchId] += 1;
        this.word_switch[docIndex][wordIndex] = newSwitchId;

        assert newTopicId > -1;
        return newTopicId;

    }

    private void build_model() {

        for (int iter = 0; iter < this.iter_num; iter++) {
            long startTime = System.currentTimeMillis();

            for (int docIndex = 0; docIndex < this.wordId_of_corpus.length; docIndex++) {
                for (int wordIndex = 0; wordIndex < this.wordId_of_corpus[docIndex].length; wordIndex++) {
                    this.topicId_of_corpus[docIndex][wordIndex] = this.sample(docIndex, wordIndex);
                }
            }

            System.out.println("finished iter :" + iter + "\tcost time:" + ((double) System.currentTimeMillis() - startTime) / 1000);
        }

    }

    private void save_result() {
        try {
            this.save_twords(20);
            this.save_theta();
            this.save_tassign();
            this.save_wordMap();
            this.save_phi();
//            this.save_background_words();
            this.save_specific_words();

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @Override
    public void run() {
        this.load_data();
        this.init_model();
        this.build_model();
        this.save_result();
    }
}
